import time

import datetime
import gzip
import json

import boto3

from framework.common.pt_logger import logger
from framework.core.crawler_result import Result


class S3Output(object):
    def __init__(self, s3_bucket, s3_folder, get_output_content=None, get_prefix_func=None, compress=True):
        self.s3_bucket = s3_bucket
        self.s3_folder = s3_folder
        self.get_output_content = get_output_content
        self.get_prefix_func = get_prefix_func
        self.compress = compress
        self.client = boto3.client('s3')

    @staticmethod
    def timestamp():
        return time.time()

    def _dump_list(self, current_task, results):
        logger.info("output List<Result>, length is %s" % len(results))
        result_sampler = results[0] if len(results) > 0 else None
        prefix = self.get_prefix_func(current_task, result_sampler)
        file_key = "%s/%s_%s.json" % (self.s3_folder, prefix, self.timestamp())
        if self.get_output_content is None:
            line = json.dumps(results, default=lambda o: o.__dict__, indent=None)
        else:
            line = self.get_output_content(current_task, results)  # JimmyMo: 这里返回的line也可能是一个数组，如何拍平统一

        if isinstance(line, list):
            logger.warning("not support output list to s3 right now")
            pass
        else:
            if self.compress:
                file_key += ".gz"
                line = gzip.compress(line.encode('utf-8'))
            response = self.client.put_object(
                Body=line,
                Bucket=self.s3_bucket,
                Key=file_key,
            )
            if response.get("ResponseMetadata", dict()).get("HTTPStatusCode", None) == 200:
                logger.info('output file to %s...' % file_key)
            else:
                logger.warning("failed to output file to %s..." % file_key)

    def _dump_single(self, current_task, result):
        prefix = self.get_prefix_func(current_task, result)
        logger.info("output is %s" % result)
        file_key = "%s/%s_%s.json" % (self.s3_folder, prefix, self.timestamp())

        if self.get_output_content is None:
            line = json.dumps(result, default=lambda o: o.__dict__, indent=None)
        else:
            line = self.get_output_content(current_task, result)  # JimmyMo: 这里返回的line也可能是一个数组，如何拍平统一
        if isinstance(line, list):
            logger.warning("not support output list to s3 right now")
            pass
        else:
            if self.compress:
                file_key += ".gz"
                line = gzip.compress(line.encode('utf-8'))
            response = self.client.put_object(
                Body=line,
                Bucket=self.s3_bucket,
                Key=file_key,
            )
            if response.get("ResponseMetadata", dict()).get("HTTPStatusCode", None) == 200:
                logger.info('output file to %s...' % file_key)
            else:
                logger.warning("failed to output file to %s..." % file_key)

    @staticmethod
    def done():
        logger.info("output done")
        pass

    def dump(self, current_task, results):
        if isinstance(results, Result):
            result = results
            self._dump_single(current_task, result)
        elif isinstance(results, list) and len(results) > 0:
            self._dump_list(current_task, results)

